{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8fb2918c-1dbd-4ff8-8ed4-4b46d97eb7ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5dbffb9e-f36f-4a5f-b3cc-a22f157926d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/root/autodl-tmp/VOCdevkit'\n",
    "VOC2007 = os.path.join(root, 'VOC2007')\n",
    "VOC2012 = os.path.join(root, 'VOC2012')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b374aa9f-e0cf-4f0b-8da6-9e2b044ce0f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 直接用VOC.yaml里的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1166dc18-86b8-470d-88f5-403e5143cfbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "name = {\n",
    "    # 手动调配了。前8个是kitti的。\n",
    "    # 有个比较麻烦的是大小写，kitti的是大写的，Voc是小写的。\n",
    "    0: 'car',\n",
    "    1: 'van',\n",
    "    2: 'truck',\n",
    "    3: 'tram',\n",
    "    4: 'person',\n",
    "    5: 'person_sitting',\n",
    "    6: 'cyclist', # 骑车和自行车还是有区别的\n",
    "    7: 'misc',\n",
    "\n",
    "    # VOC独有的\n",
    "    8: 'aeroplane',\n",
    "    9: 'bicycle',\n",
    "    10: 'bird',\n",
    "    11: 'boat',\n",
    "    12: 'bottle',\n",
    "    13: 'bus',\n",
    "    14: 'cat',\n",
    "    15: 'chair',\n",
    "    16: 'cow',\n",
    "    17: 'diningtable',\n",
    "    18: 'dog',\n",
    "    19: 'horse',\n",
    "    20: 'motorbike',\n",
    "    21: 'pottedplant',\n",
    "    22: 'sheep',\n",
    "    23: 'sofa',\n",
    "    24: 'train',\n",
    "    25: 'tvmonitor',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53dfb9f6-1d3e-4dcf-8795-1e06da38e27a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# name = {\n",
    "#   0: 'aeroplane',\n",
    "#   1: 'bicycle',\n",
    "#   2: 'bird',\n",
    "#   3: 'boat',\n",
    "#   4: 'bottle',\n",
    "#   5: 'bus',\n",
    "#   6: 'car',\n",
    "#   7: 'cat',\n",
    "#   8: 'chair',\n",
    "#   9: 'cow',\n",
    "#   10: 'diningtable',\n",
    "#   11: 'dog',\n",
    "#   12: 'horse',\n",
    "#   13: 'motorbike',\n",
    "#   14: 'person',\n",
    "#   15: 'pottedplant',\n",
    "#   16: 'sheep',\n",
    "#   17: 'sofa',\n",
    "#   18: 'train',\n",
    "#   19: 'tvmonitor',\n",
    "\n",
    "# }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d94226ae-e10b-4543-8c86-b2b51fab363b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "from tqdm import tqdm\n",
    "from utils.general import download, Path\n",
    "import shutil\n",
    "\n",
    "\n",
    "def convert_label(path, lb_path, year, image_id):\n",
    "  def convert_box(size, box):\n",
    "      dw, dh = 1. / size[0], 1. / size[1]\n",
    "      x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]\n",
    "      return x * dw, y * dh, w * dw, h * dh\n",
    "\n",
    "  in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')\n",
    "  out_file = open(lb_path, 'w')\n",
    "  tree = ET.parse(in_file)\n",
    "  root = tree.getroot()\n",
    "  size = root.find('size')\n",
    "  w = int(size.find('width').text)\n",
    "  h = int(size.find('height').text)\n",
    "\n",
    "  names = list(name.values())  # names list\n",
    "  for obj in root.iter('object'):\n",
    "      cls = obj.find('name').text\n",
    "      if cls in names and int(obj.find('difficult').text) != 1:\n",
    "          xmlbox = obj.find('bndbox')\n",
    "          bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])\n",
    "          cls_id = names.index(cls)  # class id\n",
    "          out_file.write(\" \".join([str(a) for a in (cls_id, *bb)]) + '\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "17556dcd-4ebf-4e7a-8759-04f4b73aa3bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train2012: 100%|██████████| 5717/5717 [00:02<00:00, 2348.44it/s]\n",
      "val2012: 100%|██████████| 5823/5823 [00:02<00:00, 2322.00it/s]\n",
      "train2007: 100%|██████████| 2501/2501 [00:00<00:00, 2934.59it/s]\n",
      "val2007: 100%|██████████| 2510/2510 [00:00<00:00, 2840.78it/s]\n",
      "test2007: 100%|██████████| 4952/4952 [00:02<00:00, 2464.59it/s]\n"
     ]
    }
   ],
   "source": [
    "dir = Path(root)  # dataset root dir\n",
    "\n",
    "path = dir\n",
    "for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):\n",
    "  imgs_path = dir / 'images' / f'{image_set}{year}'\n",
    "  lbs_path = dir / 'labels' / f'{image_set}{year}'\n",
    "  imgs_path.mkdir(exist_ok=True, parents=True)\n",
    "  lbs_path.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "  with open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt') as f:\n",
    "      image_ids = f.read().strip().split()\n",
    "  for id in tqdm(image_ids, desc=f'{image_set}{year}'):\n",
    "      f = path / f'VOC{year}/JPEGImages/{id}.jpg'  # old img path\n",
    "      lb_path = (lbs_path / f.name).with_suffix('.txt')  # new label path\n",
    "      #f.rename(imgs_path / f.name)  # move image\n",
    "      shutil.copy(f, imgs_path / f.name)  # 复制文件，而不是移动\n",
    "      convert_label(path, lb_path, year, id)  # convert labels to YOLO format\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5178cc71-f5cd-4947-b0e3-6f05bb93b2c8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "df7073ff-d050-4be0-9edd-116a2dc3eef8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: car\n",
      "1: van\n",
      "2: truck\n",
      "3: tram\n",
      "4: person\n",
      "5: person_sitting\n",
      "6: cyclist\n",
      "7: misc\n",
      "8: aeroplane\n",
      "9: bicycle\n",
      "10: bird\n",
      "11: boat\n",
      "12: bottle\n",
      "13: bus\n",
      "14: cat\n",
      "15: chair\n",
      "16: cow\n",
      "17: diningtable\n",
      "18: dog\n",
      "19: horse\n",
      "20: motorbike\n",
      "21: pottedplant\n",
      "22: sheep\n",
      "23: sofa\n",
      "24: train\n",
      "25: tvmonitor\n"
     ]
    }
   ],
   "source": [
    "for k, v in name.items():\n",
    "    print(f'{k}: {v}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85ac58c5-7d49-4f1f-9fdd-5a87794138aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
